import rclpy
from rclpy.node import Node
from nav_msgs.msg import OccupancyGrid
import numpy as np
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped, PointStamped
from tf2_ros import TransformListener
from tf2_ros.buffer import Buffer
from nav_msgs.msg import MapMetaData, OccupancyGrid
import heapq
import math

class MapInflation(Node):

    def __init__(self):
        super().__init__('map_inflation_node')
        self.declare_parameter('inflation_radius_value', 5)
        self.subscription = self.create_subscription(
            OccupancyGrid,
            '/map',
            self.map_callback,
            50)
        self.publisher = self.create_publisher(
            OccupancyGrid, 
            '/inflated_map', 
            10)
        self.point_sub = self.create_subscription(
            PointStamped,
            'point',
            self.point_callback,
            10
        )
        self.inflated_map=None
        self.timer = self.create_timer(1.0, self.timer_callback)
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        self.path_pub = self.create_publisher(Path, 'path', 50)
        self.map=None
        
    def timer_callback(self):
        if not self.inflated_map:
            return
        self.publisher.publish(self.inflated_map)
        self.map=np.array(self.inflated_map.data).reshape((self.inflated_map.info.height, self.inflated_map.info.width))

    def astar(self, start, end):
        open_list_fore = []
        open_list_back = []
        closed_list_fore = set()
        closed_list_back = set()
        g_fore = dict()
        g_back = dict()
        parent_fore=dict()
        parent_back=dict()
        meet=start
        heapq.heappush(open_list_fore, (0, start))
        heapq.heappush(open_list_back, (0, end))
        
        while open_list_fore and open_list_back:
            # 向前搜
            cost_fore, current_node_fore = heapq.heappop(open_list_fore)

            if current_node_fore in parent_back:
                meet=current_node_fore
                return self.extract_path(meet,parent_fore,parent_back,start,end)
            
            closed_list_fore.add(current_node_fore)
            
            for neighbor in self.get_neighbors(current_node_fore):
                if neighbor in closed_list_fore:
                    continue
                
                new_cost_fore = cost_fore + self.get_heuristic(neighbor, end)

                if neighbor not in g_fore or new_cost_fore < g_fore[neighbor]:  
                    g_fore[neighbor] = new_cost_fore
                    parent_fore[neighbor] = current_node_fore
                    heapq.heappush(open_list_fore, (new_cost_fore, neighbor))

            # 反向传播
            cost_back, current_node_back = heapq.heappop(open_list_back)

            if current_node_back in parent_fore:
                meet=current_node_back
                return self.extract_path(meet,parent_fore,parent_back,start,end)
            
            closed_list_back.add(current_node_back)
            
            for neighbor in self.get_neighbors(current_node_back):
                if neighbor in closed_list_back:
                    continue
                
                new_cost_back = cost_back + self.get_heuristic(neighbor, start)

                if neighbor not in g_back or new_cost_back < g_back[neighbor]:  
                    g_back[neighbor] = new_cost_back
                    parent_back[neighbor]=current_node_back
                    heapq.heappush(open_list_back, (new_cost_back, neighbor))
        return None
    
    def extract_path(self, s_meet,PARENT_fore,PARENT_back,start,end):
        """
        extract path from start and goal
        :param s_meet: meet point of bi-direction a*
        :return: path
        """

        # extract path for foreward part
        path_fore = [s_meet]
        s = s_meet

        while True:
            s = PARENT_fore[s]
            path_fore.append(s)
            if s == start:
                break

        # extract path for backward part
        path_back = []
        s = s_meet

        while True:
            s = PARENT_back[s]
            path_back.append(s)
            if s == end:
                break

        return list(reversed(path_fore)) + list(path_back)

    def get_neighbors(self, node):
        neighbors = []
        x, y = node
        
        directions = [(1, 0), (0, 1), (-1, 0), (0, -1)]
        
        for direction in directions:
            new_x, new_y = x + direction[0], y + direction[1]
            if 0 <= new_x < self.map.shape[1] and 0 <= new_y < self.map.shape[0] and self.map[new_y, new_x] == 0:
                neighbors.append((new_x, new_y))
        
        return neighbors


    def get_heuristic(self, node, end):
        x1, y1 = node
        x2, y2 = end
        return abs(x1 - x2) + abs(y1 - y2)

    def point_callback(self, msg):
        if self.map is None:
            return

        path = self.astar((self.world_to_map(0,0)),(self.world_to_map(3,1)))

        path_msg = Path()
        path_msg.header.stamp=self.get_clock().now().to_msg()
        path_msg.header.frame_id='map'
        for node in path:
            pose = PoseStamped()
            pose.header.stamp=self.get_clock().now().to_msg()
            pose.header.frame_id='map'
            pose.pose.position.x,pose.pose.position.y= self.map_to_world(node[0],node[1])
            path_msg.poses.append(pose)
        
        self.path_pub.publish(path_msg)


    
    def map_callback(self, msg):
        print("map is ok")
        inflation_radius_value=self.get_parameter('inflation_radius_value').get_parameter_value().integer_value
        self.inflated_map = self.inflate_map(msg, inflation_radius=inflation_radius_value) # specify the inflation radius here
        
    def inflate_map(self, map_msg, inflation_radius):
        inflated_map = OccupancyGrid()
        inflated_map.header = map_msg.header
        inflated_map.info = map_msg.info

        inflated_data = np.array(map_msg.data).reshape((map_msg.info.height, map_msg.info.width))
        
        inflated_data_copy = inflated_data.copy()

        for y in range(map_msg.info.height):
            for x in range(map_msg.info.width):
                if inflated_data[y][x] == 100:
                    for i in range(-inflation_radius, inflation_radius+1):
                        for j in range(-inflation_radius, inflation_radius+1):
                            if 0 <= y+i < map_msg.info.height and 0 <= x+j < map_msg.info.width:
                                inflated_data_copy[y+i][x+j] = 100

        inflated_list = []
        for row in inflated_data_copy:
            for val in row:
                inflated_list.append(int(val))

        inflated_map.data = inflated_list
        
        return inflated_map

    def map_to_world(self, x, y):
        mx=x*self.inflated_map.info.resolution+self.inflated_map.info.origin.position.x
        my=y*self.inflated_map.info.resolution+self.inflated_map.info.origin.position.y
        return mx, my
    
    def world_to_map(self, x, y):
        wx = x-self.inflated_map.info.origin.position.x
        wy = y-self.inflated_map.info.origin.position.y
        mx = int(wx / self.inflated_map.info.resolution)
        my = int(wy / self.inflated_map.info.resolution)
        return mx, my

def main(args=None):
    rclpy.init(args=args)

    map_inflation_node = MapInflation()

    rclpy.spin(map_inflation_node)

    map_inflation_node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
